import requests
import lzma
import pandas as pd

url_aligned = "https://data.nextstrain.org/files/ncov/open/global/aligned.fasta.xz"
url_metadata ="https://data.nextstrain.org/files/ncov/open/global/metadata.tsv.xz"

file_aligned = requests.get(url=url_aligned)
with open("aligned.fasta.xz","wb") as f:
    f.write(file_aligned.content)

file_metadata = requests.get(url=url_metadata)
with open("metadata.tsv.xz","wb") as f:
    f.write(file_metadata.content)

aliganed = dict()
with lzma.open(f"aligned.fasta.xz",'rb') as f:
    i=1
    for l in f.readlines():
        if i%2==1:
            key = l.decode("utf-8").strip(">\n")
        else:
            aliganed[key] = l.decode("utf-8").strip()
        i+=1

with lzma.open(f"metadata.tsv.xz",'rb') as f:
    with open("metadata.tsv",'wb') as f1:
        f1.write(f.read())

metadata = pd.read_csv('metadata.tsv', sep='\t')
keys = metadata.loc[metadata["country"] == "China","strain"].values.tolist()

aliganed_china = dict()
aliganed_china["strain"] = keys
aliganed_china["seq"] = list()

for k in keys:
    aliganed_china["seq"].append(aliganed[k])

aliganed_china_df = pd.DataFrame(aliganed_china)
aliganed_china_df.to_csv("reuslt.csv")
